{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Converts RGB, segmentation maps into COCO\n",
    "\n",
    "Defines the class `COCOCreator`, that creates a COCO-style json for a Replica scene. \n",
    "\n",
    "Example usage\n",
    "```\n",
    "c = CocoCreator(<root_data_dir>, <location of replica semantic json>, <list of labels>)\n",
    "c.create_coco(\n",
    "    scene='apartment_0', # replica scene name\n",
    "    coco_file_name='sample.json', # output file name\n",
    "    seg_folder='seg', # location of segmentation maps\n",
    "    pct=1 # percentage of segmentation maps to use, always 1 when creating a set on propagated labels\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sys\n",
    "if '/opt/ros/kinetic/lib/python2.7/dist-packages' in sys.path:\n",
    "    sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')\n",
    "\n",
    "import cv2\n",
    "import random\n",
    "import os\n",
    "import numpy as np\n",
    "import json\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image\n",
    "from pycococreatortools import pycococreatortools\n",
    "from IPython import embed\n",
    "from tqdm import tqdm\n",
    "from IPython.core.display import display, HTML\n",
    "import os\n",
    "import shutil\n",
    "import glob \n",
    "from detectron2.data import DatasetCatalog, MetadataCatalog\n",
    "from detectron2.data.datasets import register_coco_instances\n",
    "from detectron2.utils.visualizer import Visualizer, ColorMode\n",
    "\n",
    "# labels = ['chair', 'cushion', 'door', 'indoor-plant', 'table', 'vase']\n",
    "# labels= ['chair', 'cushion', 'door', 'lamp', 'sofa', 'table']\n",
    "labels = ['chair', 'cushion', 'door', 'indoor-plant', 'sofa', 'table']\n",
    "# Takes the output of data collection and label prop, ideally these live in the same folder called root, scene \n",
    "root_dir = '/checkpoint/apratik/finals/'\n",
    "scene = 'apartment_0'\n",
    "output_json = 'default_apt0.json'\n",
    "semantic_json_root = '/checkpoint/apratik/ActiveVision/active_vision/info_semantic'\n",
    "\n",
    "\n",
    "class CocoCreator:\n",
    "    \n",
    "    # Assumes root_data_dir has both te GT and propagated segmentation labels\n",
    "    def __init__(self, root_data_dir, semantic_json_root, labels):\n",
    "        self.rdd = root_data_dir\n",
    "        self.sjr = semantic_json_root\n",
    "        self.labels = labels\n",
    "        sorted(self.labels)\n",
    "        print(f\"Init CocoCreator for labels {self.labels}\")\n",
    "        \n",
    "    def create_coco(self, scene, coco_file_name, seg_folder, pct):\n",
    "        hsd = self.load_semantic_json(scene)\n",
    "        self.create_metadata(hsd)\n",
    "        self.create_annos(hsd, scene, seg_folder, pct)\n",
    "        self.save_json(coco_file_name)\n",
    "        self.visualize_annos(coco_file_name, scene)\n",
    "    \n",
    "    def visualize_annos(self, coco_file_name, scene):\n",
    "        DatasetCatalog.clear()\n",
    "        MetadataCatalog.clear()\n",
    "\n",
    "        register_coco_instances('foobar', {}, coco_file_name, os.path.join(self.rdd, scene, 'rgb'))\n",
    "        MetadataCatalog.get('foobar')\n",
    "        dataset_dicts = DatasetCatalog.get('foobar')\n",
    "\n",
    "        for d in dataset_dicts: #random.sample(dataset_dicts, 2):\n",
    "            img = cv2.imread(d[\"file_name\"])\n",
    "            print(d['file_name'])\n",
    "            visualizer = Visualizer(img[:, :, ::-1], metadata=MetadataCatalog.get('foobar'), scale=0.5)\n",
    "            vis = visualizer.draw_dataset_dict(d)\n",
    "            img = vis.get_image()\n",
    "            plt.figure(figsize=(12 , 8))\n",
    "            plt.imshow(img)\n",
    "            plt.show()\n",
    "        \n",
    "    def save_json(self, coco_file_name):\n",
    "        coco_output = {\n",
    "            \"info\": self.INFO,\n",
    "            \"licenses\": self.LICENSES,\n",
    "            \"categories\": self.CATEGORIES,\n",
    "            \"images\": self.IMAGES,\n",
    "            \"annotations\": self.ANNOTATIONS,\n",
    "        }\n",
    "        \n",
    "        print(f\"Dumping annotations to {coco_file_name}\")\n",
    "        with open(coco_file_name, \"w\") as output_json:\n",
    "            json.dump(coco_output, output_json)\n",
    "    \n",
    "    def create_annos(self, hsd, scene, seg, pct):\n",
    "        coco_img_id = -1\n",
    "        count = 0\n",
    "        segm_dir = os.path.join(self.rdd, scene, seg)\n",
    "        print(f\"Scene {scene}, seg dir {segm_dir}\")       \n",
    "        img_dir = os.path.join(self.rdd, scene, 'rgb')\n",
    "        fs = self.get_segm_files(segm_dir, pct)\n",
    "        print(f\"Creating COCO annotations for {len(fs)} images \\n img_dir {img_dir}\")\n",
    "        \n",
    "        for f in tqdm(fs):\n",
    "            image_id = int(f.split('.')[0])\n",
    "            try:\n",
    "                prop_path = os.path.join(segm_dir, \"{:05d}.npy\".format(image_id))\n",
    "                annot = np.load(prop_path).astype(np.uint32)\n",
    "            except Exception as e:\n",
    "                print(e)\n",
    "                continue\n",
    "\n",
    "\n",
    "            img_filename = \"{:05d}.jpg\".format(image_id)            \n",
    "            img = Image.open(os.path.join(img_dir, img_filename))\n",
    "\n",
    "            # COCO ID and names\n",
    "            coco_img_id += 1\n",
    "\n",
    "            image_info = pycococreatortools.create_image_info(\n",
    "                coco_img_id, os.path.basename(img_filename), img.size\n",
    "            )\n",
    "\n",
    "            self.IMAGES.append(image_info)\n",
    "        \n",
    "            # for each annotation add to coco format\n",
    "            for i in np.sort(np.unique(annot.reshape(-1), axis=0)):\n",
    "                try:\n",
    "                    if hsd[\"id_to_label\"][i] < 1 or self.label_id_dict[hsd[\"id_to_label\"][i]] not in labels:# or hsd[\"id_to_label\"][i] not in self.label_id_dict:\n",
    "                        continue\n",
    "                    category_info = {\"id\": self.new_old_id[hsd[\"id_to_label\"][i]], \"is_crowd\": False}\n",
    "#                     print(f'category_info {category_info}')\n",
    "                except:\n",
    "                    continue\n",
    "\n",
    "                binary_mask = (annot == i).astype(np.uint32)\n",
    "\n",
    "                annotation_info = pycococreatortools.create_annotation_info(\n",
    "                    count, coco_img_id, category_info, binary_mask, img.size, tolerance=2\n",
    "                )\n",
    "                if annotation_info is not None:\n",
    "                    self.ANNOTATIONS.append(annotation_info)\n",
    "                    count += 1\n",
    "        \n",
    "        \n",
    "    def create_metadata(self, hsd):\n",
    "        self.INFO = {}\n",
    "        self.LICENSES = [{}]\n",
    "        self.CATEGORIES = []\n",
    "        self.IMAGES = []\n",
    "        self.ANNOTATIONS = []\n",
    "        \n",
    "        self.label_id_dict = {}\n",
    "        self.new_old_id = {}\n",
    "        idc = 1\n",
    "        for obj_cls in hsd[\"classes\"]:\n",
    "            if obj_cls[\"name\"] in self.labels:\n",
    "                self.CATEGORIES.append({\"id\": idc, \"name\": obj_cls[\"name\"], \"supercategory\": \"shape\"})\n",
    "                self.label_id_dict[obj_cls[\"id\"]] = obj_cls[\"name\"]\n",
    "                self.new_old_id[obj_cls['id']] = idc\n",
    "                idc += 1\n",
    "        \n",
    "        \n",
    "    def load_semantic_json(self, scene):\n",
    "        replica_root = '/datasets01/replica/061819/18_scenes'\n",
    "        habitat_semantic_json = os.path.join(replica_root, scene, 'habitat', 'info_semantic.json')\n",
    "        print(f\"Using habitat semantic json {habitat_semantic_json}\")\n",
    "        with open(habitat_semantic_json, \"r\") as f:\n",
    "            hsd = json.load(f)\n",
    "        if hsd is None:\n",
    "            print(\"Semantic json not found!\")\n",
    "        return hsd\n",
    "    \n",
    "    def get_segm_files(self, segm_dir, pct):\n",
    "        cs = [os.path.basename(x) for x in glob.glob(os.path.join(segm_dir, '*.npy'))]\n",
    "        cs.sort()\n",
    "        frq = 1/pct\n",
    "        fs = []\n",
    "        for x in range(0, len(cs), int(frq)):\n",
    "            fs.append(cs[x])\n",
    "        return fs    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seg_folder=[f'gt{x}p2_rand_{y}' for x in range(5,30,5) for y in range(3)]\n",
    "print(seg_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# seg_folder=[f'gt10p{x}_h1' for x in range(2, 10, 2)]\n",
    "\n",
    "# Create the active and baseline jsons for each setting\n",
    "for seg_folder in seg_folder:\n",
    "    cactive = CocoCreator('/checkpoint/apratik/finals/straightline', semantic_json_root, labels)\n",
    "    cactive.create_coco(\n",
    "        scene='apartment_0', \n",
    "        coco_file_name='straightline_apt0_'+seg_folder+'.json', \n",
    "        seg_folder='pred_'+seg_folder, \n",
    "        pct=1 # always 1 for training on propagated labels\n",
    "    )\n",
    "\n",
    "    cbase = CocoCreator('/checkpoint/apratik/finals/default', semantic_json_root, labels)\n",
    "    cbase.create_coco(\n",
    "        scene='apartment_0', \n",
    "        coco_file_name='default_apt0_'+seg_folder+'.json', \n",
    "        seg_folder='pred_'+seg_folder, \n",
    "        pct=1 # always 1 for training on propagated labels\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# c.create_coco('apartment_0', 'default_apt0_100gt.json', 'pred_label_gt100', 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a Test set\n",
    "\n",
    "# ctest = CocoCreator('/checkpoint/apratik/ActiveVision/active_vision/replica_random_exploration_data', semantic_json_root, labels)\n",
    "# ctest.create_coco('frl_apartment_1', 'frlapt1_20n.json', 'seg', 0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "loco",
   "language": "python",
   "name": "loco"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
